import csv
import json
import optparse
import os
import pickle
import time
from functools import reduce

def mkdir(path):
    folder = os.path.exists(os.getcwd() + '/' + path)
    if not folder:  # 判断是否存在文件夹如果不存在则创建为文件夹
        os.makedirs(os.getcwd() + '/' + path)  # makedirs 创建文件时如果路径不存在会创建这个路径


def ReadFile(path):
    # file = open('/mnt/c/Users/Administrator/Desktop/ezTaxon_16s.human-gut.fasta')
    # file = open('C:/Users/Administrator/Desktop/ezTaxon_16s.human-gut.fasta')
    # file = open('C:/Users/Administrator/Desktop/kmer.fasta')
    file = open(path)
    dict = {}
    for line in file:
        if line.startswith('>'):
            name = line.split('>')[1].split(' ')[0]+'\n'
            dict[name] = ''
        else:
            dict[name] += line.replace('\n', '')
    return dict
    # dict = {}
    # for record in SeqIO.parse(path, "fasta"):
    #     dict[record.id] = record.seq
    # return dict


def Kmerseq(K):
    arr = reduce(lambda x, y: [i + j for i in x for j in y], [['A', 'T', 'C', 'G']] * K)
    # dict = {}
    # for seq in arr:
    #     dict[seq] = seq
    return arr


def getSeqDict(str, k, oldArray, indexDict):

    for i in range(0, len(str) - k + 1):
        t = str[i:i + k]
        if t in indexDict:
            index = indexDict[t]
            oldArray[index] += 1
    return oldArray
def GetKmerFile(path,K,dir):
    name = os.path.dirname(dir)
    # 获取Kmer相关文件
    # 输入的fasta文件
    mkdir(name)
    file = ReadFile(path)
    # print(len(file))
    kmer = Kmerseq(K)
    indexDict = {}
    for i in range(len(kmer)):
        indexDict[kmer[i]] = i
    rf = open(dir, 'w', encoding='utf-8')
    # rf = open('ezTaxon_16s.csv', 'w', encoding='utf-8')
    rf.write(",".join(kmer))
    rf.write("\n")
    count = 0
    for line in file:
        count += 1
        oldArray = [0] * len(kmer)
        vs = getSeqDict(file[line], K, oldArray, indexDict)
        vs.append(line)
        vs = map(str, vs)
        newLine = ",".join(vs)
        rf.write(newLine)
        # if count % 100 == 0:
        #     print(count)
        #     print(time.time() - beforetime)
    rf.close()

if __name__ == '__main__':
    before = time.time()
    parser = optparse.OptionParser()
    parser.add_option('--in', help='[File path], fasta file, required.', dest='infile', action='store',
                      default='ezTaxon_16s.human-gut.fasta')
    parser.add_option('--kmer', help='[Int], K-mer length, default 8.', dest='kmer', action='store', default='8')
    # parser.add_option('--class', help='[File path], group file, format: SeqID<TAB>Label, required.', dest='classfile',
    #                   action='store', default='group.csv')
    # parser.add_option('--model',
    #                   help='[String], trainning model, must be one of ["NB","SVC","RF","GB","DT","RENB","RESVR","RERF","REGB"], default NB.',
    #                   dest='modelname', action='store', default='SVC')
    # parser.add_option('--out_model', help='[File path], output trained model file, required.', dest='modelfile',
    #                   action='store', default='')
    parser.add_option('--out_dir',
                      help='[Directory path], output directory, contains model summary information, required.',
                      dest='dirfile', action='store', default='dir')
    # parser.add_option('--fp', help='[opionts], parameter setting information for each model.', dest='parameter_file',
    #                   action='store', default='')
    opts, args = parser.parse_args()
    data = eval(str(opts))
    path = data['infile']
    K = int(data['kmer'])
    # classfile = data['classfile']
    # modelfile = data['modelfile']
    dir = data['dirfile']
    # model = data['modelname']
    # paramfile = data['parameter_file']
    # if paramfile!='':
    #     paramtxt = open(paramfile)
    # else:
    #     paramtxt = ''
        # feature = data['feature']
    GetKmerFile(path, K, dir)